#include "stl-helper.h"
#include "gene_processing.h"


/******************* GIBBS sampler ****************************/

MotifGibbsSampler::MotifGibbsSampler(boost::program_options::variables_map& vm, const Sequences& seqs)
  : vm_(vm), seqs_(seqs), acgt("ACGT")

{
  outputFile_ = vm["output"].as<std::string>();
  freqFile_ = vm["freq"].as<std::string>();
  debugFile_ = vm.count("debug") ? vm["debug"].as<std::string>() : outputFile_ + ".debug";
  max_iter_ = vm["iter"].as<int>();
  motif_length_ = vm["mtlen"].as<int>();
}

void MotifGibbsSampler::run() 
{
  std::ofstream debug(debugFile_.c_str());
  debug.close();

  prepare_count();
  prev_idx_ = -2;
  // Initialize motif positions in sequences
  motif_position_ = vi_t(sz(seqs_), 0);
  for (int i = 0; i < sz(seqs_); i++)
    motif_position_[i] = rand() % (sz(seqs_[i]) - motif_length_ + 1);

  for (int iter = 0; iter < max_iter_; iter++) {
    std::cerr << iter << "\t "; std::cerr.flush(); //MSG("iter = " << iter);
    if (iter % 10 == 9) std::cerr << ENDL;
    for (int i = 0; i < sz(seqs_); i++) {
      gibbs_sampler(i);
    }
    print_debug(iter);
    
    if (iter % 1000 == 999) {
      phase_shift();
      
      std::string motif(motif_length_, 'x');
      for (int c = 0; c < 4; c++) {
	std::cerr << acgt[c];
	for (int pos = 0; pos <= motif_length_; pos++) { 
	  bool ok = true;
	  for (int c1 = 0; c1 < 4; c1++) 
	    if (freq_[pos][c1] > freq_[pos][c]) ok = false;
	  std::cerr << "\t " << (ok?"+":"") << freq_[pos][c];
	  if (pos > 0 && ok) motif[pos-1] = acgt[c];
	}
	std::cerr << ENDL;
      }
      std::cerr << "motif = " << motif << ENDL
		<< "likelihood = " << log_likelihood() << ENDL;
    }

  } // for iter
  
  compute_max_position();

  prev_idx_ = -2;
  count(-1);
  compute_freq(-1);
}

void MotifGibbsSampler::compute_max_position()
{
  prev_idx_ = -2;
  count(-1);
  compute_freq(-1);

  // compute Maximum Likelihood Ratio position
  for (int i = 0; i < sz(seqs_); i++) {
    vd_t likelihood;
    compute_likelihood(i, likelihood);
    int max_pos = 0;
    for (int pos = 0; pos < sz(likelihood); pos++) {
      if (likelihood(pos) > likelihood(max_pos))
	max_pos = pos;
    }
    motif_position_[i] = max_pos;
  }
}

void MotifGibbsSampler::print_debug(int iter)
{
  count(-1);
  compute_freq(-1);

  std::ofstream debug(debugFile_.c_str(), std::ios::app);
  debug << "iter = " << iter << ENDL;
  for (int c = 0; c < 4; c++) {
    debug << acgt[c];
    for (int pos = 0; pos <= motif_length_; pos++) { 
      bool ok = true;
      for (int c1 = 0; c1 < 4; c1++) 
	if (freq_[pos][c1] > freq_[pos][c]) ok = false;
      debug << "\t " << (ok?"+":"") << freq_[pos][c];
    }
    debug << ENDL;
  }
  debug.close();
}

void MotifGibbsSampler::prepare_count()
{
  int N = sz(seqs_);
  pre_count_ = std::vector<vi_t> (N, vi_t(4, 0));

  for (int i = 0; i < N; i++) {
    const std::string& s_i = seqs_[i];
    for (int pos = 0; pos < sz(s_i); pos++) {
      char c = s_i[pos];
      int c_pos = acgt.find(c);
      pre_count_[i][c_pos] ++;
    }
  }
}

void MotifGibbsSampler::gibbs_sampler(int seq_idx)
{
  count(seq_idx);
  compute_freq(seq_idx);
  sample(seq_idx);
}

void MotifGibbsSampler::count(int seq_idx)
{
  if (prev_idx_ == -2) { // count for the first time
    count_ = vvi_t(motif_length_ + 1, vi_t(4, 0));
    for (int i = 0; i < sz(seqs_); i++) 
      if (i != seq_idx) {
	const std::string& s_i = seqs_[i];
	int motif_pos = motif_position_[i];
	
	for (int c_pos = 0; c_pos < 4; c_pos++)
	  count_[0][c_pos] += pre_count_[i][c_pos];
	
	for (int pos = motif_pos; pos < motif_pos + motif_length_; pos++) {
	  char c = s_i[pos];
	  int c_pos = acgt.find(c);
	  count_[0][c_pos]--; 
	  count_[pos-motif_pos+1][c_pos]++;
	}
      }
  }
  else { // from the second time, do something smart
    if (prev_idx_ >= 0) { // add back the previously excluded sequence
      const std::string& s_i = seqs_[prev_idx_];
      int motif_pos = motif_position_[prev_idx_];

      for (int c_pos = 0; c_pos < 4; c_pos++)
	count_[0][c_pos] += pre_count_[prev_idx_][c_pos];
      
      for (int pos = motif_pos; pos < motif_pos + motif_length_; pos++) {
	char c = s_i[pos];
	int c_pos = acgt.find(c);
	count_[0][c_pos]--; 
	count_[pos-motif_pos+1][c_pos]++;
      }
    }

    if (seq_idx >= 0) { // remove the current sampling sequence from counting
      const std::string& s_i = seqs_[seq_idx];
      int motif_pos = motif_position_[seq_idx];
      for (int c_pos = 0; c_pos < 4; c_pos++)
	count_[0][c_pos] -= pre_count_[seq_idx][c_pos];
      for (int pos = motif_pos; pos < motif_pos + motif_length_; pos++) {
	char c = s_i[pos];
	int c_pos = acgt.find(c);
	count_[0][c_pos]++; 
	count_[pos-motif_pos+1][c_pos]--;
      }
    }    
  }

  // MSG("seq_idx = " << seq_idx << " prev_idx = " << prev_idx_);
  // for (int c = 0; c < 4; c++) {
  //   std::cerr << acgt[c];
  //   for (int pos = 0; pos <= motif_length_; pos++)
  //     std::cerr << "\t " << count_[pos][c];
  //   std::cerr << ENDL;
  // }

  prev_idx_ = seq_idx;
}

// void MotifGibbsSampler::count(int seq_idx)
// {
//   count_ = vvi_t(motif_length_ + 1, vi_t(4, 0));

//   for (int i = 0; i < sz(seqs_); i++) 
//     if (i != seq_idx) {
//       const std::string& s_i = seqs_[i];
//       int motif_pos = motif_position_[i];
//       for (int pos = 0; pos < sz(s_i); pos++) {
// 	char c = s_i[pos];
// 	int c_pos = acgt.find(c);
// 	if (pos < motif_pos || pos >= motif_pos + motif_length_) { // background
// 	  count_[0][c_pos]++;
// 	}
// 	else {  // motif
// 	  count_[pos-motif_pos+1][c_pos]++;
// 	}
//       }
//     }
// }

void MotifGibbsSampler::compute_freq(int seq_idx)
{
  double b = 0.5, B = b*4; // pseudo-count
  freq_ = vvd_t(motif_length_+1, arma::zeros<vd_t>(4));
  double total_background = 0;
  int N = sz(seqs_);
  for (int i = 0; i < N; i++)
    total_background += (i != seq_idx) ? sz(seqs_[i]) - motif_length_ : 0;

  for (int c = 0; c < 4; c++) {
    freq_[0][c] = (count_[0][c] + b) / (total_background + B);
    for (int pos = 1; pos <= motif_length_; pos++) 
      freq_[pos][c] = (count_[pos][c] + b) / (N - (seq_idx == -1 ? 0 : 1) + B);
  }
}

void MotifGibbsSampler::sample(int seq_idx)
{
  vd_t likelihood;
  compute_likelihood(seq_idx, likelihood);
  
  double sum = arma::sum(likelihood);
  double r = arma::randu() * sum, csum = 0;
  int pos;
  int len = sz(seqs_[seq_idx]) - motif_length_ + 1;
  for (pos = 0; pos < len; pos++) {
    csum += likelihood(pos);
    if (csum >= r) break;
  }

  motif_position_[seq_idx] = pos < len ? pos : len - 1;
}

void MotifGibbsSampler::compute_likelihood(int seq_idx, vd_t& likelihood)
{
  const std::string& s_i = seqs_[seq_idx];
  int len = sz(s_i) - motif_length_ + 1;
  likelihood = arma::ones<vd_t>(len);
  
  for (int pos = 0; pos < len; pos++) { // compute likelihood in each position
    for (int i = 0; i < motif_length_; i++) {
      char c = s_i[pos+i];
      int c_pos = acgt.find(c);
      likelihood(pos) *= freq_[i+1][c_pos] / freq_[0][c_pos];
    }
  }
}

void MotifGibbsSampler::phase_shift()
{
  compute_max_position();

  vi_t current_pos = motif_position_;

  int dpos[5] = {-2, -1, 0, 1, 2};
  double log_likelihood[5];
  //double sum = LogNumber::zero();
  int max_d = 0;
  for (int d = 0; d < 5; d++) {
    shift(current_pos, motif_position_, dpos[d]);
    prev_idx_ = -2;
    count(-1);
    compute_freq(-1);
    
    double likelihood = this->log_likelihood();

    log_likelihood[d] = likelihood / sz(seqs_);
    if (log_likelihood[d] > log_likelihood[max_d])
      max_d = d;
    //sum = LogNumber::add(sum, log_likelihood[d]);
  }

  //double r = LogNumber::randu(sum); // random [0, e^sum]
  double csum = LogNumber::zero();
  int d = 0;
  //MSG(ENDL << r);
  for (d = 0; d < 5; d++) {
    csum = LogNumber::add(csum, log_likelihood[d]);
    std::cerr << "\t " << log_likelihood[d] << " (" << csum << ")";
    //if (csum >= r) break;
  }
  
  d = dpos[max_d]; //d < 5 ? dpos[d] : dpos[4];
  shift(current_pos, motif_position_, d);
  
  prev_idx_ = -2;
  count(-1);
  compute_freq(-1);

  MSG(ENDL << "Phase shift d = " << d);
}

void MotifGibbsSampler::shift(const vi_t& current_pos, vi_t& pos, int dpos)
{
  int N = sz(seqs_);
  for (int i = 0; i < N; i++) {
    int p = current_pos[i] + dpos;
    int len = sz(seqs_[i]) - motif_length_ + 1;
    if (p < 0) p = 0;
    if (p >= len) p = len - 1;
    pos[i] = p;
  }
}

double MotifGibbsSampler::log_likelihood()
{
  double likelihood = 0;
  for (int i = 0; i < sz(seqs_); i++) {
    const std::string& s_i = seqs_[i];
    int pos = motif_position_[i];
    for (int j = 0; j < motif_length_; j++) {
      char c = s_i[pos+j];
      int c_pos = acgt.find(c);
      likelihood += log(freq_[j+1][c_pos]) - log(freq_[0][c_pos]);
    }
  }
  return likelihood;
}
